//
// Created by chizuru on 2025/9/18.
//

#pragma once

#include <shared_mutex>

namespace aop {

// -----------------------------
// Advisable<R, Args...>
// -----------------------------
    template<typename R, typename... Args>
    class Advisable {
    public:
        using FuncType = std::function<R(Args...)>;
        using Before   = std::function<void(Args&...)>;
        using AfterValue = std::function<void(R&, Args&...)>;
        using AfterVoid  = std::function<void(Args&...)>;
        using ProceedFn  = std::function<R(Args&...)>;
        using Around   = std::function<R(ProceedFn, Args&...)>;

        using token_t = uint64_t;

        Advisable() : original(nullptr), token_counter_(1) {}
        explicit Advisable(FuncType f) : original(std::move(f)), token_counter_(1) {}

        // set original target later
        void set_original(FuncType f) { std::unique_lock lk(mu_); original = std::move(f); }

        // add_before: accepts any callable convertible to Before (i.e. void(Args&...))
        template<typename F>
        token_t add_before(F &&f) {
            Before bf = std::function<void(Args&...)>(std::forward<F>(f));
            return add_impl(befores_, std::move(bf));
        }

        template<typename F>
        token_t add_after_value(F &&f) {
            AfterValue af = std::function<void(R&, Args&...)>(std::forward<F>(f));
            return add_impl(afters_value_, std::move(af));
        }

        template<typename F>
        token_t add_after_void(F &&f) {
            AfterVoid av = std::function<void(Args&...)>(std::forward<F>(f));
            return add_impl(afters_void_, std::move(av));
        }

        template<typename F>
        token_t add_around(F &&f) {
            Around a = std::function<R(ProceedFn, Args&...)>(std::forward<F>(f));
            return add_impl(arounds_, std::move(a));
        }

        bool remove(token_t t) {
            std::unique_lock lk(mu_);
            if (remove_from_vec(befores_, t)) return true;
            if (remove_from_vec(afters_value_, t)) return true;
            if (remove_from_vec(afters_void_, t)) return true;
            if (remove_from_vec(arounds_, t)) return true;
            return false;
        }

        // call operator: uses "storage + reference" approach
        R operator()(Args... args) {
            // pack args into decay_t storage (we let advices operate on these internal copies by reference)
            std::tuple<std::decay_t<Args>...> storage(std::forward<Args>(args)...);

            // snapshot advices under shared lock
            std::vector<Before> befores_copy;
            std::vector<AfterValue> afters_value_copy;
            std::vector<AfterVoid> afters_void_copy;
            std::vector<Around> arounds_copy;

            {
                std::shared_lock lk(mu_);
                befores_copy.reserve(befores_.size());
                for (auto &p : befores_) befores_copy.push_back(p.second);
                afters_value_copy.reserve(afters_value_.size());
                for (auto &p : afters_value_) afters_value_copy.push_back(p.second);
                afters_void_copy.reserve(afters_void_.size());
                for (auto &p : afters_void_) afters_void_copy.push_back(p.second);
                arounds_copy.reserve(arounds_.size());
                for (auto &p : arounds_) arounds_copy.push_back(p.second);
            }

            // helper to apply befores/afters on tuple storage (elts are lvalue refs to elements)
            auto call_befores = [&](auto &tpl){
                std::apply([&](auto &...elts){
                    for (auto &b : befores_copy) b(elts...);
                }, tpl);
            };
            auto call_afters_value = [&](R &res, auto &tpl){
                std::apply([&](auto &...elts){
                    for (auto &a : afters_value_copy) a(res, elts...);
                }, tpl);
            };
            auto call_afters_void = [&](auto &tpl){
                std::apply([&](auto &...elts){
                    for (auto &a : afters_void_copy) a(elts...);
                }, tpl);
            };

            // run befores
            call_befores(storage);

            // prepare base proceed that calls original
            ProceedFn base = [this](Args&... as)->R {
                // original might be null if not set; throw in that case
                if (!original) throw std::bad_function_call();
                if constexpr (std::is_void_v<R>) {
                    original(as...);
                    return;
                } else {
                    return original(as...);
                }
            };

            // compose around chain (reverse order)
            ProceedFn proceed = base;
            for (auto it = arounds_copy.rbegin(); it != arounds_copy.rend(); ++it) {
                Around around = *it;
                ProceedFn prev = proceed;
                proceed = [around, prev](Args&... as)->R {
                    return around(prev, as...);
                };
            }

            // invoke proceed with references to storage elements
            if constexpr (std::is_void_v<R>) {
                std::apply([&](auto &...elts){ proceed(elts...); }, storage);
                call_afters_void(storage);
                return;
            } else {
                R result = std::apply([&](auto &...elts)->R { return proceed(elts...); }, storage);
                call_afters_value(result, storage);
                return result;
            }
        }

    private:
        // helper impls
        template<typename Vec, typename Fn>
        token_t add_impl(Vec &vec, Fn &&fn) {
            std::unique_lock lk(mu_);
            token_t t = token_counter_.fetch_add(1, std::memory_order_relaxed);
            vec.emplace_back(t, std::forward<Fn>(fn));
            return t;
        }

        template<typename Vec>
        bool remove_from_vec(Vec &vec, token_t t) {
            for (auto it = vec.begin(); it != vec.end(); ++it) {
                if (it->first == t) { vec.erase(it); return true; }
            }
            return false;
        }

        // members
        FuncType original;
        std::shared_mutex mu_;
        std::atomic<token_t> token_counter_;

        // advices containers: pair<token, fn>
        std::vector<std::pair<token_t, Before>> befores_;
        std::vector<std::pair<token_t, AfterValue>> afters_value_;
        std::vector<std::pair<token_t, AfterVoid>> afters_void_;
        std::vector<std::pair<token_t, Around>> arounds_;
    };

    // helper factory: create from std::function
    template<typename R, typename... Args>
    inline Advisable<R, Args...> make_advisable(std::function<R(Args...)> f) {
        return Advisable<R, Args...>(std::move(f));
    }

    // bind member (raw pointer)
    template<typename Obj, typename Rm, typename... Argsm>
    inline Advisable<Rm, Argsm...> bind_member(Obj* obj, Rm(Obj::*mf)(Argsm...)) {
        auto f = [obj, mf](Argsm... args)->Rm { return (obj->*mf)(std::forward<Argsm>(args)...); };
        return make_advisable<Rm, Argsm...>(std::function<Rm(Argsm...)>(std::move(f)));
    }

    template<typename Obj, typename Rm, typename... Argsm>
    inline Advisable<Rm, Argsm...> bind_member(const Obj* obj, Rm(Obj::*mf)(Argsm...) const) {
        auto f = [obj, mf](Argsm... args)->Rm { return (obj->*mf)(std::forward<Argsm>(args)...); };
        return make_advisable<Rm, Argsm...>(std::function<Rm(Argsm...)>(std::move(f)));
    }

    // bind member with weak_ptr safety
    template<typename Obj, typename Rm, typename... Argsm>
    inline Advisable<Rm, Argsm...> bind_member_weak(std::weak_ptr<Obj> wp, Rm(Obj::*mf)(Argsm...)) {
        auto f = [wp, mf](Argsm... args)->Rm {
            if (auto sp = wp.lock()) {
                return (sp.get()->*mf)(std::forward<Argsm>(args)...);
            }
            throw std::bad_function_call();
        };
        return make_advisable<Rm, Argsm...>(std::function<Rm(Argsm...)>(std::move(f)));
    }

} // namespace aop